# -*- coding: utf-8 -*-
# EvaluationData.py
from surprise.model_selection import train_test_split
from surprise.model_selection import LeaveOneOut
from surprise import KNNBaseline


class EvaluationData:
    """
    用于准备和管理推荐系统评估所需的数据。
    """

    def __init__(self, data, popularityRankings):
        """
        初始化 EvaluationData 类。

        :param data: 原始数据集，通常是一个 Surprise 数据集。
        :param popularityRankings: 项目的流行度排名信息，用于后续计算覆盖率等指标。
        """
        self.rankings = popularityRankings

        # 使用整个数据集构建一个完整的训练集，用于评估模型的整体性能。
        self.fullTrainSet = data.build_full_trainset()
        # 包含用户尚未评分的所有项目，用于生成 Top-N 推荐。
        self.fullAntiTestSet = self.fullTrainSet.build_anti_testset()

        # 将数据集划分为训练集（75%）和测试集（25%），用于评估模型的预测准确性。
        self.trainSet, self.testSet = train_test_split(data, test_size=0.25, random_state=1)

        # 每个用户仅保留一条评分作为测试集，其余评分作为训练集。
        LOOCV = LeaveOneOut(n_splits=1, random_state=1)
        for train, test in LOOCV.split(data):
            self.LOOCVTrain = train
            self.LOOCVTest = test

        # 包含用户在 Leave-One-Out 测试集中尚未评分的所有项目。
        self.LOOCVAntiTestSet = self.LOOCVTrain.build_anti_testset()

        # 使用余弦相似度计算项目之间的相似度矩阵，用于衡量推荐结果的多样性。
        sim_options = {
            'name': 'cosine',
            'user_based': False,
            'shrinkage': 100  # 收缩参数避免除零
        }

        # KNN 计算相似度矩阵
        self.simsAlgo = KNNBaseline(sim_options=sim_options)
        self.simsAlgo.fit(self.fullTrainSet)

    def GetFullTrainSet(self):
        """
        返回完整的训练集，用于评估模型的整体性能。

        :return: 完整的训练集。
        """
        return self.fullTrainSet

    def GetFullAntiTestSet(self):
        """
        返回用户尚未评分的所有项目集合。

        :return: 用户尚未评分的所有项目集合。
        """
        return self.fullAntiTestSet

    def GetAntiTestSetForUser(self, userID: int):
        """
        获得某用户的未评分的测试数据集
        :param userID: 用户ID
        :return:
        """
        trainset = self.fullTrainSet
        try:
            # 将原始用户ID转换为内部用户ID
            u = trainset.to_inner_uid(str(userID))
        except ValueError:
            print(f"用户 {userID} 不存在于训练集中，将返回空推荐列表")
            return []

        # 用于填充未评分项目
        fill = trainset.global_mean
        anti_testset = []
        # 用户评分过的数据
        user_items = set([j for (j, _) in trainset.ur[u]])
        # 添加类型转换确保兼容性
        anti_testset += [
            (str(trainset.to_raw_uid(u)),  # 确保用户ID字符串类型
             str(trainset.to_raw_iid(i)),  # 确保物品ID字符串类型
             fill)
            for i in trainset.all_items() if i not in user_items
        ]  # [原始用户ID, 原始项目ID, 填充值]
        return anti_testset

    def GetTrainSet(self):
        """
        返回 75/25 划分的训练集。

        :return: 训练集。
        """
        return self.trainSet

    def GetTestSet(self):
        """
        返回 75/25 划分的测试集。

        :return: 测试集。
        """
        return self.testSet

    def GetLOOCVTrainSet(self):
        """
        返回 Leave-One-Out 划分的训练集。

        :return: Leave-One-Out 训练集。
        """
        return self.LOOCVTrain

    def GetLOOCVTestSet(self):
        """
        返回 Leave-One-Out 划分的测试集。

        :return: Leave-One-Out 测试集。
        """
        return self.LOOCVTest

    def GetLOOCVAntiTestSet(self):
        """
        返回 Leave-One-Out 划分的 AntiTestSet。

        :return: Leave-One-Out AntiTestSet。
        """
        return self.LOOCVAntiTestSet

    def GetSimilarities(self):
        """
        返回项目之间的相似度矩阵。

        :return: 项目之间的相似度矩阵。
        """
        return self.simsAlgo

    def GetPopularityRankings(self):
        """
        返回项目的流行度排名信息。

        :return: 项目的流行度排名信息。
        """
        return self.rankings
